import os
from utils.ExamplePostProcess import post_process_yolo_v5
import kp
import numpy as np
import cv2
from .base_tpu import TPU


class PredictYolov5(TPU):
    def __init__(self):
        super().__init__()
        # 提取文件名, 输出对应的文件路径
        file_name = os.path.basename(__file__)
        print(f"{file_name}.py path : {self.pwd}")

    def predict(self, img, imgsz, conf=None, iou=None, model_nth=0):
        results_box = []
        if conf or iou:
            if imgsz[0] > 640 or imgsz[1] > 480:
                img = cv2.resize(img, (imgsz[0], imgsz[1]))
            else:
                imgsz = [640, 480]

            img = cv2.cvtColor(src=img, code=cv2.COLOR_BGR2BGR565)
            out_boxes = self.after_poscessor_yolov5(self.device_group, self.model_nef_descriptor,
                                                    model_nth, img, conf, iou)

            for o in out_boxes.box_list:
                x1 = int(o.x1 * 640 / imgsz[0])
                y1 = int(o.y1 * 480 / imgsz[1])
                x2 = int(o.x2 * 640 / imgsz[0])
                y2 = int(o.y2 * 480 / imgsz[1])
                label = int(o.class_num)
                score = float(o.score)
                if x2 > 640:
                    x2 = 640
                if x1 < 0:
                    x1 = 0
                if y2 > 480:
                    y2 = 480
                if y1 < 0:
                    y1 = 0

                w = x2 - x1
                h = y2 - y1
                x = (x1 + w / 2 - 1)
                y = (y1 + h / 2 - 1)
                # 将处理后的坐标、宽高和标签存储到列表中
                results_box.append([int(x), int(y), int(w), int(h), score, label])
            return results_box
        else:
            img = cv2.resize(img, (imgsz[0], imgsz[1]))
            img = cv2.cvtColor(src=img, code=cv2.COLOR_BGR2BGR565)
            label = self.after_poscessor_yolov8(self.device_group, self.model_nef_descriptor,
                                                model_nth, img)
            return label


    def after_poscessor_yolov5(self, device_group, model_nef_descriptor, model_nth, img, thresh_value, iou_value):

        generic_inference_input_descriptor = kp.GenericImageInferenceDescriptor(
            model_id=model_nef_descriptor.models[model_nth].id,
            inference_number=0,
            input_node_image_list=[
                kp.GenericInputNodeImage(
                    image=img,
                    image_format=kp.ImageFormat.KP_IMAGE_FORMAT_RGB565,
                    resize_mode=kp.ResizeMode.KP_RESIZE_ENABLE,
                    padding_mode=kp.PaddingMode.KP_PADDING_CORNER,
                    normalize_mode=kp.NormalizeMode.KP_NORMALIZE_YOLO
                )
            ]
        )

        try:
            kp.inference.generic_image_inference_send(device_group=device_group,
                                                        generic_inference_input_descriptor=generic_inference_input_descriptor)
            generic_raw_result = kp.inference.generic_image_inference_receive(device_group=device_group)
        except kp.ApiKPException as exception:
            print(' - Error: inference failed, error = {}'.format(exception))
            exit(0)

        inf_node_output_list = []

        for node_idx in range(generic_raw_result.header.num_output_node):
            inference_float_node_output = kp.inference.generic_inference_retrieve_float_node(node_idx=node_idx,
                                                                                            generic_raw_result=generic_raw_result,
                                                                                            channels_ordering=kp.ChannelOrdering.KP_CHANNEL_ORDERING_CHW
                                                                                            )
            inf_node_output_list.append(inference_float_node_output)

        yolo_result = post_process_yolo_v5(inference_float_node_output_list=inf_node_output_list,
                                        hardware_preproc_info=generic_raw_result.header.hw_pre_proc_info_list[0],
                                        thresh_value=thresh_value, iou_value = iou_value)

        return yolo_result

    # def after_poscessor_yolov8(self, device_group, model_nef_descriptor, model_nth, img):

    #     generic_inference_input_descriptor = kp.GenericImageInferenceDescriptor(
    #         model_id=model_nef_descriptor.models[model_nth].id,
    #         inference_number=0,
    #         input_node_image_list=[
    #             kp.GenericInputNodeImage(
    #                 image=img,
    #                 image_format=kp.ImageFormat.KP_IMAGE_FORMAT_RGB565,
    #                 resize_mode=kp.ResizeMode.KP_RESIZE_ENABLE,
    #                 padding_mode=kp.PaddingMode.KP_PADDING_CORNER,
    #                 normalize_mode=kp.NormalizeMode.KP_NORMALIZE_YOLO

    #             )
    #         ]
    #     )

    #     try:
    #         kp.inference.generic_image_inference_send(device_group=device_group,
    #                                                     generic_inference_input_descriptor=generic_inference_input_descriptor)
    #         generic_raw_result = kp.inference.generic_image_inference_receive(device_group=device_group)
    #     except kp.ApiKPException as exception:
    #         print(' - Error: inference failed, error = {}'.format(exception))
    #         exit(0)

    #     inf_node_output_list = []

    #     for node_idx in range(generic_raw_result.header.num_output_node):
    #         inference_float_node_output = kp.inference.generic_inference_retrieve_float_node(node_idx=node_idx,
    #                                                                                         generic_raw_result=generic_raw_result,
    #                                                                                         channels_ordering=kp.ChannelOrdering.KP_CHANNEL_ORDERING_CHW
    #                                                                                         )
    #         inf_node_output_list.append(inference_float_node_output)

    #     tmp = []

    #     for o_n in inf_node_output_list:
    #         o_array = o_n.ndarray.copy()
    #         tmp.append(o_array)

    #     res = self.postprocess(tmp)

    #     return np.argmax(res)

    # def softmax(self, A):
    #     e = np.exp(A)
    #     return e / np.sum(e, keepdims=True)

    # def postprocess(self, pre_output):
    #     score = self.softmax(pre_output)
    #     return score

#tpu = PredictYolov5()
#tpu.link()
#tpu.model_init("/opt/SONGKE_BOT/app/ARM4/static_grab/model/", True)
#tpu.release()
